classdef etsCounter
    %etsCounter Counterfactual simulations of emissions market outcomes
    %
    %   An etsCounter object takes an etsEstim object that has estimates of
    %   an abatement model. It then 
        
    properties
        
        % Objects created on initialization
        estim;   % Estimates of the model
        plants;  % Plant-level data, inherited from estimation
        bids;    % Bid-level data, inherited from estimation
        sample;  % Indicator for counterfactual sample from plant-period table
                
        ip;    % Indices of plants in obj.plants, with length of obj.bids
        ib;    % Indices of bids offered by a plant in obj.bids
        
        S = 20; % Number of simulation draws from estimated plant distribution
        
        %% Specify counterfactuals to be run
        %    'regime' -- One of {'Market','Command'}
        %       Market  : Emissions trading system (treatment)
        %       Command : Command & control regime (control)
        %    'input'   -- One of {'Emissions','Cost'}
        %    'ermodel' -- Model for status quo emissions rate.  One of 
        %       'flat','flat_eps','heat','heat_eps','heat_eps_rho'
        %       flat     : Constant emissions rate for all firms
        %       flat_eps : Constant emissions rate, epsilon error drawn
        %       heat     : Emissions rate conditional on Hi
        %       heat_eps : Emissions rate conditional on Hi, epsilon
        %         error drawn independently of Xi_it
        %       heat_eps_rho : Emissions rate conditional on Hi, epsilon
        %         error drawn with correlation with Xi_it
        %   'cost' -- Which model of plant-period costs to use
        %       iso     : Iso-elastic cost
        %       step    : Single Step Function Cost
        scenario = struct(        ...
            'regime',  'Market',    ...
            'input',   'Emissions', ...
            'ermodel', 'flat',  ...
            'rho',     -0.1, ...
            'costs',   'iso',...
            'iso_type', 'no_hetero');
        noisy = true;
        
        % Parameters governing when and how market is to be cleared
        periods; 
        period;
        
        cap = [ 100 : 100 : 2000 ];
        level;
        costs; 

        % Should we allow a free cyclone when calculating step function
        % MACs
        free_cyc_step = true;
        % Should we include minimum period usage as step function minimum
        include_usage_min = true;
                     
        %% Results from counterfactual runs
       
        % Simulation draws of plant X period effects
        
        % Cell array of tables of counterfactual emissions and cost outcomes
        %   One array entry for each scenario
        beta;
        output_it; 
        output_t;
        
        output_itk;
        output_tk;
        
        % Outcomes
        outVars    = {'Eit','Zit','MACit'};
        outVarLabs = {'Emissions','Abatement cost','Marginal cost'};
    end
        
    properties (Dependent, SetAccess = private)
        
        % Indices into aggregate summary table (period-level)
        idx_agg_t;  % Index into period
        idx_agg_tk; % Index into period-level pair
        
        % Indices into plant summary table (plant-period-level)
        idx_out_t;  % Index into period, for all plants
        idx_out_tk; % Index into period-level, for all plants
    end
    
    methods
        %% Initialization of counterfactual object
        function obj = etsCounter( etsEstim, emission_caps )
            % Initialize etsCounter object
            
            % Store data on plants and bids
            obj.estim   = etsEstim;
            obj.plants  = etsEstim.plants;
            obj.bids    = removevars(etsEstim.bids,...
                [etsEstim.fe_pp,etsEstim.fe_plants,etsEstim.fe_periods]);
            
            % Sample limitations: Drop control plants
            if ~ismember('treatment',obj.plants.Properties.VariableNames)
                obj.plants.treatment = ones(size(obj.plants.id_period));
            end
            obj.sample = obj.plants.treatment==1 & ...
                ~isnan(obj.plants.id_period);
 
            % Periods and emissions levels for which to solve
            obj.periods  = unique(obj.plants.id_period(obj.sample));
            
            % Initialize output: plant X period level balanced panel
            obj.output_it = obj.plants(obj.sample,{'id_plant','id_period',...
                'Hi','Ebari','ERit'});
            planttab = tabulate(obj.output_it.id_period);
            assert(all(planttab(:,2)==156));
            obj.output_it.shutdown = (obj.plants.Eit(obj.sample)==0);
            
            % Initialize output: period aggregates
            obj.output_t    = table(obj.periods,'VariableNames',{'id_period'});
            obj.output_t.Pt = nan(length(obj.periods),1);

            % Expand output tables to have entries for each emissions level
            if nargin > 1
                obj.cap = emission_caps;
            end
            
        end
        
        function [ output_itk, output_tk ] = expandLevels( obj )
            % Expand output tables for each level of emissions
            output_itk = expandTable( obj.output_it, obj.cap' );
            output_tk  = expandTable( obj.output_t, obj.cap' );           
        end
        
        function [ obj ] = importModelEstimates( obj )
            % Bring relevant fixed effects estimates from estimated model 
            % into output data set, to be used to compute counterfactuals
               
            % set appropriate model_num for use in importing depending on
            % what iso_type is set
            if obj.scenario.iso_type == "no_hetero"
                model_num = 3;
            elseif obj.scenario.iso_type == "hetero"
                model_num = 4;
            end

            % Extract abatement elasticity
            betaIdx = find(contains(obj.estim.fits{model_num}.Coefficients.Name,'logE'));
            beta_coef = obj.estim.fits{model_num}.Coefficients.Estimate(betaIdx);
            
            n_plants = length(unique(obj.output_it.id_plant));
            if model_num == 3
                obj.beta = beta_coef * ones(n_plants, 1);
                % obj.beta = beta_coef;
            elseif model_num == 4
                beta_i = zeros(n_plants, 1);
                plant_apcd = obj.bids(:, ["id_plant","apcd_max"]);
                plant_apcd = unique(plant_apcd);
                for i = 1:n_plants
                    type = plant_apcd(plant_apcd.id_plant==i,"apcd_max").Variables;
                    if type > 0
                        if type > 3
                            type = 3;
                        end
                        if type == 2
                            type = 1;
                        end
                        if type > 1
                            type = 2;
                        end
                        beta_i(i) = beta_coef(type);
                        % beta_i(i) = median(beta_coef);
                    else
                        beta_i(i) = median(beta_coef);
                    end
                end
                obj.beta = beta_i;
            end
            
            % Extract plant X period effects
            ppIdx     = contains(obj.estim.fits{model_num}.Coefficients.Name,'id_pp');
            xi_it_hat = obj.estim.fits{model_num}.Coefficients.Estimate(ppIdx);
            pp_name   = obj.estim.fits{model_num}.Coefficients.Name(ppIdx);
            pp_num    = cellfun(@(x) str2num(extractAfter(x,'id_pp_')),pp_name);
                   
            % Set plant period averages from their fixed effect in plain
            % fixed effect regression without emissions
            pp_avg_p = obj.estim.fits{5}.Coefficients.Estimate(ppIdx);

            % Create period and plant identifiers from coefficient names
            id_period = floor(pp_num/10000);
            id_plant  = mod(pp_num,1000);
            fes       = table( id_period, id_plant, xi_it_hat, pp_avg_p );
            
            % Merge period X plant effects back into output data frame
            obj.output_it = outerjoin(obj.output_it,fes,...
                'Keys',{'id_period','id_plant'},...
                'MergeKeys',true);

        end      
        
        function [ obj ] = imputeMissingCostParameters( obj )
            % Impute unobserved cost parameter at the plant-period level
            %   for firms that did not bid or had missing data
            
            % xi_it
            % Impute xi_it for plant-periods where the plant did not have
            %   an estimated xi_it but *does* have a xi_it for at least one
            %   other compliance period
            xi_plant = grpstats(obj.output_it(:,{'id_plant','xi_it_hat'}),...
                {'id_plant'},{'mean'});
            xi_plant = removevars(xi_plant,'GroupCount');
            obj.output_it = outerjoin(obj.output_it,xi_plant,...
                'Keys','id_plant','MergeKeys',true);
            xi_it_na = isnan(obj.output_it.xi_it_hat);
            obj.output_it.xi_it_hat( xi_it_na ) = ...
                obj.output_it.mean_xi_it_hat( xi_it_na );
            obj.output_it = removevars(obj.output_it,'mean_xi_it_hat');
            obj.output_it = sortrows(obj.output_it,{'id_period','id_plant'});
            obj.output_it = movevars(obj.output_it,'id_period','Before','id_plant');
            
            % Impute xi_it for plant-periods where the plant does not have
            %   an estimated xi_it in any period
            xi_it_na = isnan(obj.output_it.xi_it_hat);
            xi_agg_median = median( obj.output_it.xi_it_hat(~xi_it_na) );
            obj.output_it.xi_it_hat( xi_it_na ) = xi_agg_median;
            assert( ~any( isnan(obj.output_it.xi_it_hat) ) );

            % pp_avg_p
            % Impute pp_avg_p for plant-periods where the plant did not have
            %   an estimated pp_avg_p but *does* have a pp_avg_p for at least one
            %   other compliance period
            pp_plant = grpstats(obj.output_it(:,{'id_plant','pp_avg_p'}),...
                {'id_plant'},{'mean'});
            pp_plant = removevars(pp_plant,'GroupCount');
            obj.output_it = outerjoin(obj.output_it,pp_plant,...
                'Keys','id_plant','MergeKeys',true);
            pp_avg_p_na = isnan(obj.output_it.pp_avg_p);
            obj.output_it.pp_avg_p( pp_avg_p_na ) = ...
                obj.output_it.mean_pp_avg_p( pp_avg_p_na );
            obj.output_it = removevars(obj.output_it,'mean_pp_avg_p');
            obj.output_it = sortrows(obj.output_it,{'id_period','id_plant'});
            obj.output_it = movevars(obj.output_it,'id_period','Before','id_plant');
            

            % Impute pp_avg_p for plant-periods where the plant does not have
            %   an estimated pp_avg_p in any period
            pp_avg_p_na = isnan(obj.output_it.pp_avg_p);
            pp_avg_p_agg_median = median( obj.output_it.pp_avg_p(~pp_avg_p_na) );
            obj.output_it.pp_avg_p( pp_avg_p_na ) = pp_avg_p_agg_median;
            assert( ~any( isnan(obj.output_it.pp_avg_p) ) );
        end
        
        %% Retrieve indices into period-level and plant-period-level tables
        %    These indices are used to select a subset of observations in
        %    the output tables to update. The obj.output_itk table contains
        %    plant-period-level data for different levels of emissions. The
        %    obj.output_tk contains period-level data for different levels
        %    of emissions.
        
        % Index into period, all levels of emissions
        function idx_agg_t = get.idx_agg_t( obj )
            idx_agg_t = ...
                ( obj.output_tk.id_period == obj.period );
        end
        
        % Index into period-level pair
        function idx_agg_tk = get.idx_agg_tk( obj )
            idx_agg_tk = ...
                ( obj.output_tk.id_period == obj.period ) & ...
                ( obj.output_tk.id_level  == obj.level  );
        end
        
        % Index into period, for all plants and levels
        function idx_out_t = get.idx_out_t( obj )
            idx_out_t = ...
                ( obj.output_itk.id_period == obj.period );
        end
        
        % Index into period-level, for all plants
        function idx_out_tk = get.idx_out_tk( obj )
            idx_out_tk = ...
                ( obj.output_itk.id_period == obj.period ) & ...
                ( obj.output_itk.id_level  == obj.level  );
        end
        
        %% Run counterfactuals
        function obj = run( obj )
            % Run counterfactuals for all scenarios provided
            rng(0);

            % Import coefficient estimates from model estimates
            obj = obj.importModelEstimates;
            
            % Impute xi_it for plant-periods where they are not estimated
            obj = obj.imputeMissingCostParameters;
            
            [ obj.output_itk, obj.output_tk ] = expandLevels( obj );

            if obj.noisy
                    header = [ repmat('*',1,80) '\n' ];
                    fprintf(1,header);
                    fprintf(1,'Counterfactual scenario\n');
                    fprintf(1,'\tRegime: %s\n',obj.scenario.regime);
                    if strcmp(obj.scenario.regime,'Command')
                        fprintf(1,'\tEmissions rate: %s\n',obj.scenario.ermodel);
                    end
                    fprintf(1,'\tCost Functional Form: %s\n', obj.scenario.costs);
                    if strcmp(obj.scenario.costs,'iso')
                        fprintf(1,'\tIso-elastic parameter estimation model: PP FEs + %s\n',obj.scenario.iso_type);
                    end
                    fprintf(1,'\tInput:  %s\n\n',obj.scenario.input);
            end

            % Add columns to output_it which are useful in step function
            % MAC curve regime processing
            if obj.scenario.costs == "step"
                obj = obj.stepAugment;
            end
            
            switch obj.scenario.regime
                case 'Market'
                    if obj.scenario.costs == "iso"
                        obj = obj.solveMarketAll;
                    elseif obj.scenario.costs == "step"
                        obj = obj.solveStepAll;
                    end
                case 'Command'
                    obj = obj.solveCommandAll;                    
            end
        end
        
        % Add columns to output_it which are useful in step function MAC
        % curve regime processing
        function [ obj ] = stepAugment( obj )
          
            min_emissions = groupfilter(obj.plants, "id_plant", @(x) x == min(x), "Eit");
            min_emissions = renamevars(min_emissions, "Eit", "minPeriodEmissions");
            min_emissions = groupfilter(min_emissions, "id_plant", @(x) x == min(x), "id_period");
            min_emissions = min_emissions(:, ...
            ["id_plant" "apcd_max" "minPeriodEmissions" "D_cyc" "D_scr" "D_bf" "D_esp"]);
 
            obj.output_it = outerjoin(obj.output_it, min_emissions, ...
                        'Keys',{'id_plant'},...
                        'Type', 'Left',...
                        'MergeKeys',true);

            obj.output_it.E_noabate = obj.output_it.Ebari;

            % If we allow free cyc usage then the max possible
            % abatement for those with scrubber is already abating 80%
            if obj.free_cyc_step
                obj.output_it.E_noabate(obj.output_it.D_cyc==1) = .2 * ...
                    obj.output_it.E_noabate(obj.output_it.D_cyc==1);
            end


            % For apcd_max =0 or missing set their abate emissions equal to
            % no abate emissions
            obj.output_it.E_abate = obj.output_it.E_noabate;
            

            % If free cyclone then for those with only cyclone already have
            % their abatement included so don't further abate
            if obj.free_cyc_step
                obj.output_it.E_abate(obj.output_it.apcd_max==1) = obj.output_it.E_noabate(obj.output_it.apcd_max==1);
            else
                obj.output_it.E_abate(obj.output_it.apcd_max==1) = .2 * obj.output_it.E_noabate(obj.output_it.apcd_max==1);
            end

            obj.output_it.E_abate(obj.output_it.apcd_max==2) = .01 * obj.output_it.E_noabate(obj.output_it.apcd_max==2);
            obj.output_it.E_abate(obj.output_it.apcd_max==3) = .06 * obj.output_it.E_noabate(obj.output_it.apcd_max==3);
            obj.output_it.E_abate(obj.output_it.apcd_max==4) = .003 * obj.output_it.E_noabate(obj.output_it.apcd_max==4);

            if obj.include_usage_min
                obj.output_it.E_abate = min(obj.output_it.E_abate, obj.output_it.minPeriodEmissions);
            end

            % Add dummies for abatement tech and columns for levels of
            % emissions with and without abatement for step function
            obj.output_itk = outerjoin(obj.output_itk, obj.output_it(:,...
                    ["id_period" "id_plant" "minPeriodEmissions" "apcd_max" "D_cyc" "D_scr" "D_bf" "D_esp" "E_abate" "E_noabate"]), ...
                 "Keys", ["id_period" "id_plant"], ...
                'Type', 'Left',...
                'MergeKeys',true);
            
        end


        %% Solve for Market results when using a step function approach
        function [ obj ] = solveStepAll( obj )
            
            % Solve for emissions market equilibrium in all periods
            for t = obj.periods'
                %Update compliance period, which determines what period
                %   the market equilibrium will be solved *for*
                obj.period = t;

                for k = obj.cap
                    obj.level = k;

                    output_itk_working = obj.output_itk(obj.idx_out_tk,:);

                    % Solve for equilibrium price and emissions levels
                    Et = Inf;
                    price=1;
                    while (Et > obj.level) & (price<100)
                        price = price + .01;
                        Eit = output_itk_working.E_noabate;
                        will_abate = output_itk_working.pp_avg_p <= price;
                        Eit(will_abate) = output_itk_working.E_abate(will_abate);
                        Et = sum(Eit);
                    end

                    % Assign equilibrium price to the relevant period
                    obj.output_tk.Pt( obj.idx_agg_tk ) = price ;

                    % Assign plant emissions to the relevant period
                    obj.output_itk.Eit( obj.idx_out_tk ) = Eit;

                    % Assign aggregate emissions to the relevant period
                    obj.output_tk.Et( obj.idx_agg_tk ) = Et;

                    % Assign Price
                    Zit = zeros(length(Eit),1);
                    Zit(will_abate) = output_itk_working.pp_avg_p(will_abate) .* ...
                        (output_itk_working.E_noabate(will_abate) - output_itk_working.E_abate(will_abate));
                    Zt = sum(Zit);

                    obj.output_itk.Zit( obj.idx_out_tk ) = Zit;
                    obj.output_tk.Zt( obj.idx_agg_tk ) = Zt;
                end
            end
            
            obj.output_itk = outerjoin(obj.output_itk, obj.output_tk,...
                "Type","left",...
                "MergeKeys",true);
        end

        %% Counterfactuals: emissions trading market equilibrium
        function [ obj ] = solveMarketAll( obj )
            % Solve for emissions market equilibrium in all periods
            for t = obj.periods'
                
                % Update compliance period, which determines what period
                %   the market equilibrium will be solved *for*
                obj.period = t;

                % Emissions levels for which to clear the market
                for k = obj.cap
                    obj.level = k;
                    
                    % Market-clearing price
                    obj.output_tk.Pt( obj.idx_agg_tk ) = obj.solveMarketOne;

                    % Assign plant emissions
                    obj = obj.assignPlantEmissions;
                    
                    % Calculate total cost of abatement
                    [ Zt, Zit ] = obj.solveTotalCostsWrapper;
                    obj.output_itk.Zit( obj.idx_out_tk ) = Zit;
                    obj.output_tk.Zt( obj.idx_agg_tk ) = Zt;
                end
            end
        end
         
        function [ Pstar ] = solveMarketOne( obj )
            % Solve for emissions market equilibrium in one period given:
            %    - Plant characteristics (observable)
            %    - Plant characteristics (unobservable)
            %    - Permit supply
            % return market-clearing price and aggregate emissions.
     
            % Parameters for market-clearing
            P0   = [ 1e-5 1e8 ];
            Qcap = obj.level;
            
            % Find market-clearing price (QD(P*) - QS = 0)
            objfun = @( P ) (obj.solvePlantEmissions( P ) - Qcap);
            Pstar  = fzero( objfun, P0 );     
        end
        
        function [ Et, Eit ] = solvePlantEmissions( obj, P )
            % Solve for plant-level emissions in a compliance period
            %   as a function of market prices.
            % Return: 
            %   Et  - Aggregate emissions 
            %   Eit - Vector of plant-period emissions
            
            % Indices of all plants active in period-level
            t = ( obj.idx_out_tk );
            
            % Calculate emissions
            E1   = exp( -( obj.output_itk.xi_it_hat(t) ) );
            % E2   = obj.plant_t.Hi.^(-beta2);
            E2   = ones(size(E1));
            EPit = P.^(1./obj.beta);

            Eit = EPit .* [ E1 .* E2 ].^(1./obj.beta); % Plant-period emissions
            Et  = sum(Eit);                           % Aggregate emissions
        end
        
        function [ obj ] = assignPlantEmissions( obj )
            % Assign plant emissions in permit market equilibrium
            
            % Solve emissions for each plant
            [ Eit, Eit_equil ] = ...
                obj.solvePlantEmissions( obj.output_tk.Pt(obj.idx_agg_tk) );
            
            % Assign plant emissions to the relevant period
            obj.output_itk.Eit( obj.idx_out_tk ) = Eit_equil;
            
            % Assign aggregate emissions to the relevant period
            obj.output_tk.Et( obj.idx_agg_tk ) = Eit;
        end

        %% Counterfactuals: Command-and-control status quo regime
        function [ obj ] = solveCommandAll( obj )
            
            for t = obj.periods'
                
                % Update compliance period, which determines what period
                %   the market equilibrium will be solved *for*
                obj.period = t;

                % Emissions levels for which to clear the market
                for k = obj.cap
                    obj.level = k;
                    
                    % Assign plant emissions
                    [ Ets, Eits ] = obj.commandPlantEmissions;
                    obj.output_itk.Eit( obj.idx_out_tk ) = mean(Eits,2);
                    obj.output_tk.Et( obj.idx_agg_tk )   = mean(Ets,2);
                    
                    % Plants can't be assigned to less than their minimum level
                    % possible under the step function regime
                    if obj.scenario.costs == "step"
                        obj.output_itk.Eit( obj.idx_out_tk ) = ...
                            max(obj.output_itk.Eit( obj.idx_out_tk ),...
                            obj.output_itk.E_abate( obj.idx_out_tk ));
                        obj.output_tk.Et( obj.idx_agg_tk ) = sum(obj.output_itk.Eit( obj.idx_out_tk ));
                    end
                    
                    % Calculate total cost of abatement
                    if obj.scenario.costs == "iso"
                        [ Zts, Zits ] = obj.solveTotalCostsWrapper( Eits );
                        obj.output_itk.Zit( obj.idx_out_tk ) = mean(Zits,2);
                        obj.output_tk.Zt( obj.idx_agg_tk )   = mean(Zts,2);
                    elseif obj.scenario.costs == "step"
                        obj = obj.solveCommandStepTotalCosts;
                    end
                    
                    % Shadow cost of marginal abatement
                    obj.output_tk.Pt( obj.idx_agg_tk ) = obj.shadowCost;
                end
            end
        end
        
        function [ Ets, Eits ] = commandPlantEmissions( obj, Et )
            % Assign plant emissions in command-and-control regime
            %   as a function of aggregate emissions level.
            % Input: 
            %   Et - Aggregate emissions level
            % If not passed, then aggregate emissions level taken from
            % obj.level property of the counterfactual object.
            str = RandStream('mt19937ar','Seed',20220524); % Random # stream
            
            % Indices of all plants active in period-level
            t  = obj.period;
            ti = ( obj.idx_out_tk );
            N  = length(obj.output_itk.Hi(ti));

            % Assign plant characteristics needed to calculate emissions 
            Hi  = obj.output_itk.Hi(ti);
            obj.output_itk.logHi(ti) = log(obj.output_itk.Hi(ti));
            if nargin < 2
                Et  = obj.level;
            end
            
            % Draw emissions rates
            switch obj.scenario.ermodel
                
                % Constant emissions rate R for all firms
                case 'flat'
                    logERit = predict(obj.estim.cfit{t,1},obj.output_itk(ti,:));
                    ERit    = exp(logERit);
                    ERits   = ERit(:,ones(1,obj.S));
                    
                % Constant emissions rate R_i for all firms, with error
                case 'flat_eps'
                    logERit = predict(obj.estim.cfit{t,1},obj.output_itk(ti,:));
                    sigEps  = sqrt(obj.estim.cfit{t,1}.MSE);
                    ERits   = exp(logERit(:,ones(1,obj.S)) + ...
                                normrnd(0,sigEps,N,obj.S));
                    
%                     muER  = obj.estim.parER{obj.period,1};
%                     sigER = obj.estim.parER{obj.period,2};
%                     N     = length(obj.output_itk.Hi(ti));
%                     ERits = lognrnd( muER, sigER, N, obj.S );
                    
                % Emissions rate a function of heat output
                case 'heat'
                    logERit = predict(obj.estim.cfit{t,2},obj.output_itk(ti,:));
                    ERit    = exp(logERit);
                    ERits   = ERit(:,ones(1,obj.S));
                    
                % Emissions rate a function of heat output, with error
                case 'heat_eps'
                    logERit = predict(obj.estim.cfit{t,2},obj.output_itk(ti,:));
                    sigEps  = sqrt(obj.estim.cfit{t,1}.MSE);
                    N       = length(obj.output_itk.Hi(ti));
                    ERits   = exp(logERit(:,ones(1,obj.S)) + ...
                                normrnd(0,sigEps,N,obj.S));
                    
                % Emissions rate a function of heat output, with error
                %   Error correlated with the plant's marginal cost shock
                case 'heat_eps_rho'
                    logERit = predict(obj.estim.cfit{t,2},obj.output_itk(ti,:));
                    sigEps  = sqrt(obj.estim.cfit{t,1}.MSE);
                    N       = length(obj.output_itk.Hi(ti));
                    
                    % Draw R shock correlated with marginal cost xi_it_hat
                    xi_it = obj.output_itk.xi_it_hat(ti);
                    alpha = obj.scenario.rho * sigEps / std(xi_it);
                    sigZ  = sigEps * sqrt( 1 - obj.scenario.rho^2 );
                    Z1    = normrnd(0,sigZ,N,obj.S);
                    Z2    = alpha * (xi_it - mean(xi_it));
                    eps_its = Z1 + Z2(:,ones(1,obj.S));
                    ERits   = exp(logERit(:,ones(1,obj.S)) + eps_its);        
            end
            
            % Solve for plant emissions by scaling emissions factor
            %   distribution up or down proportionally to agg. emissions
            Eits_raw = ERits .* Hi(:,ones(1,obj.S)); 
            Efactor = Et(:,ones(1,obj.S)) ./ sum(Eits_raw,1);
            Eits    = Efactor(ones(N,1),:) .* Eits_raw; % Re-scaled emissions    
            Ets     = sum(Eits,1);
        end
        
        function [ Pstar ] = shadowCost( obj )
            % Solve for shadow cost of abatement in one period given:
            %    - Plant characteristics (observable)
            %    - Plant characteristics (unobservable)
            %    - Permit supply
            % return shadow cost and aggregate emissions.
            
            % Total cost of abatement at Et = obj.level 
            [ Zt, Zit ] = obj.solveTotalCostsWrapper;
            obj.output_itk.Zit( obj.idx_out_tk )    = Zit;
            obj.output_tk.Zt( obj.idx_agg_tk ) = Zt;
            
            % Total cost of abatement at Et - delta
            delta = 2;
            Et    = obj.level - delta;
            
            %   Plant emissions
            [ ~, Eit ] = obj.commandPlantEmissions( Et );
            
            %   Costs given emissions
            t     = ( obj.idx_out_tk ); % Index based on obj period and level 
            xi_it = obj.output_itk.xi_it_hat(t);
            Ebari = obj.output_itk.Ebari(t);
            betaZ = obj.beta;
            [ Zt_delta ] = solveTotalCosts( xi_it, Ebari, Eit, betaZ );
                        
            % Shadow cost of marginal emissions reduction
            Pstar = (Zt_delta - Zt) / delta;
        end
        
        function [ obj ] = solveCommandStepTotalCosts( obj )

            output_itk_working = obj.output_itk( obj.idx_out_tk ,:);

            % If required emissions are above max then no abatement
            Zit = zeros(length(output_itk_working.Eit), 1);

            % If required emissions are below min then abate until the minimum
            temp_index = output_itk_working.Eit < output_itk_working.E_abate;
            Zit(temp_index) = ...
                 (output_itk_working.E_noabate(temp_index) - ...
                 output_itk_working.E_abate(temp_index)) .* ...
                 output_itk_working.pp_avg_p(temp_index);

            % If required emissions are between max and min then abate until
            % necessary
            temp_index = (output_itk_working.E_abate <= output_itk_working.Eit) & ...
                (output_itk_working.Eit <= output_itk_working.E_noabate);
            Zit(temp_index) = ...
                (output_itk_working.E_noabate(temp_index) - ...
                output_itk_working.Eit(temp_index)) .* ...
                output_itk_working.pp_avg_p(temp_index);

            obj.output_itk.Zit( obj.idx_out_tk ) = Zit;
            obj.output_tk.Zt( obj.idx_agg_tk )   = sum(Zit);

        end


        %% Cost calculations
        function [ Zts, Zits ] = solveTotalCostsWrapper( obj, Eits )
            % Calculate total costs given emissions

            % Indices of all plants active in period-level
            t = ( obj.idx_out_tk ); % Index based on obj period and level

            % Assign plant characteristics needed to calculate costs
            xi_it = obj.output_itk.xi_it_hat(t);
            Ebari = obj.output_itk.Ebari(t);
            if nargin < 2
                Eits   = obj.output_itk.Eit(t);
            end

            % Solve for variable abatement costs
            betaZ = obj.beta; % Elasticity of MAC wrt emissions
            [ Zts, Zits ] = solveTotalCosts( Eits, xi_it, Ebari, betaZ );

        end
        
        function [ Pt, cap ] = invertMarketPrice( obj, Zt )
           % Calculate the market price and quantity that yield a given
           %   aggregate cost of pollution abatement
            
        end

         % Evaluate abatement cost elasticity at control emissions
         function [ ] = abateCostElasticityControlEmissions ( obj, period_num )
             output_tk = obj.output_tk(obj.output_tk.id_period==period_num,{'Et','Zt'});
             E_ref     = 2.4e5;
             delta     = 0.2e5;

             del_Z = (output_tk.Zt(output_tk.Et==E_ref+delta) - ...
                 output_tk.Zt(output_tk.Et==E_ref-delta)) ./ output_tk.Zt(output_tk.Et==E_ref);
             del_E = 2*delta / E_ref;
             elas_ZE   = del_Z / del_E;
             fprintf('Elasticity of abatement cost with respect to emissions: %3.2f\n',...
                 elas_ZE);
         end
        
        %% Additional methods of etsCounter object in external files
        
        % Tabulate outcomes of counterfactuals
        [ ] = tabulate( obj, filename, drop, title, label, footnote, ...
            addmetarows );
        
        % Scatter plot of prices over compliance periods
        % scatter(ctr.output_agg.id_period,ctr.output_agg.Pt)
        
       
    end % Methods
    
end % Classdef

%% Helper functions called by etsCounter methods
function [ tableLong ] = expandTable( table, levels )
    % Expand table by creating a new entry for each value of the variable
    %   levels.  For example, if the input table is:
    %      id     xvar
    %      1       2.3
    %      2       1.2
    %  And levels = [ 2 3 4 ]', then the output will be
    %      id      level       xvar
    %      1           2        2.3
    %      1           3        2.3
    %      1           4        2.3
    %      2           2        1.2
    %      2           3        1.2
    %      2           4        1.2
    T = height(table);
    K = length(levels);
    tableLong = table(repelem([1:T]',K),:);
    id_level  = kron(ones(T,1),levels);
    tableLong = addvars(tableLong,id_level);
end

function [ Zts, Zits ] = solveTotalCosts( Eits, xi_it, Ebari, betaZ )
    % Calculate total costs given emissions
    %   Can accept as input either:
    %   -- Eit  [N x 1] vector for each plant
    %   -- Eits [N x S] matrix for plant emissions in each of S simulations
    S = size(Eits,2);
    
    % Plant X simulation level total costs
    Z1 = exp( xi_it ) .* ( 1 ./ (betaZ + 1) );
    Z2 = Ebari.^(betaZ+1);
    Z3 = Eits.^(betaZ+1);
    Zits = Z1(:,ones(1,S)) .* (Z2(:,ones(1,S)) - Z3);

    % Aggregate total costs across simulations
    Zts = sum(Zits,1);    
end


